import requests
import base64
import json
from typing import Optional


class BaiduOCR:
    def __init__(self, api_key: str, secret_key: str):
        self.api_key = api_key
        self.secret_key = secret_key
        self.access_token = self._get_access_token()

    def _get_access_token(self) -> Optional[str]:
        url = "https://aip.baidubce.com/oauth/2.0/token"
        params = {
            "grant_type": "client_credentials",
            "client_id": self.api_key,
            "client_secret": self.secret_key
        }
        response = requests.post(url, params=params)
        if response.status_code == 200:
            return response.json().get("access_token")
        else:
            print(f"Error getting access token: {response.text}")
            return None

    def recognize_text(self, image_path: str) -> Optional[dict]:
        if not self.access_token:
            print("Access token is not available.")
            return None

        url = f"https://aip.baidubce.com/rest/2.0/ocr/v1/general_basic?access_token={self.access_token}"

        with open(image_path, "rb") as image_file:
            img = base64.b64encode(image_file.read())

        headers = {'content-type': 'application/x-www-form-urlencoded'}
        payload = {"image": img}
        response = requests.post(url, data=payload, headers=headers)

        if response.status_code == 200:
            return response.json()
        else:
            print(f"Error recognizing text: {response.text}")
            return None


def main():
    # 替换为您自己的API密钥和秘密密钥
    api_key = "ReYneqYKuZlE2wJcnz58nIdJ"
    secret_key = "rhkJw5OSDGuAUv9zJEbPMcHg7qeAbtOo"

    ocr = BaiduOCR(api_key, secret_key)

    # 替换为您的图像的路径
    image_path = "uploads/屏幕截图 2023-09-27 084326.png"

    result = ocr.recognize_text(image_path)
    if result:
        print(json.dumps(result, ensure_ascii=False, indent=2))
    else:
        print("Failed to recognize text.")


if __name__ == "__main__":
    main()